/* Copyright (C) 2011 clueless <clueless.mmopeon@gmail.com>
 *
 * This file is part of MMO Peon.
 *
 * MMO Peon is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * MMO Peon is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MMO Peon.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdlib.h>
#include <windows.h>
#include <stdint.h>
#include <string.h>
#include <stddef.h>

#define DO_INJECT           0
#define DO_EJECT            1
#define DO_CALL             2
#define PROCESS_INJECT_DLL  PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | \
                            PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ

// TODO: - lots of error checking
//       - Freeing up resources

/*
CPU Disasm
Address   Hex dump          Command
00000000  90            NOP
00000001  90            NOP
00000002  68 EFBEADDE   PUSH DEADBEEF       Push address of DLL path
00000007  BA EFBEADDE   MOV EDX,DEADBEEF    move address of GetModuleHandle in EDX
0000000C  FFD2          CALL EDX            call it
0000000E  68 EFBEADDE   PUSH DEADBEEF       push address of Function name
00000013  50            PUSH EAX            push return value of GetModuleHandle
00000014  BA EFBEADDE   MOV EDX,DEADBEEF    grab the address of GetProcAddress
00000019  FFD2          CALL EDX            call it...
0000001B  FFD0          CALL EAX            and call the address we just got returned (lol if errors happen enjoy your crash)
0000001D  C20400        RET 0x04
00000020  00                                Filler
*/

// EB FE
#define CALL_STUB       "\x90\x90"              \
                        "\x68\xEF\xBE\xAD\xDE"  \
                        "\xBA\xEF\xBE\xAD\xDE"  \
                        "\xFF\xD2"              \
                        "\x68\xEF\xBE\xAD\xDE"  \
                        "\x50"                  \
                        "\xBA\xEF\xBE\xAD\xDE"  \
                        "\xFF\xD2"              \
                        "\xFF\xD0"              \
                        "\xC2\x04\x00"          \
                        "\x00"
struct CallStub
{
    uint8_t     res1_[3];
    uintptr_t   dllString;
    uint8_t     res2_;
    uintptr_t   getModuleHandle;
    uint8_t     res3_[3];
    uintptr_t   funcString;
    uint8_t     res4_[2];
    uintptr_t   getProcAddress;
    uint8_t     res5_[7];
    char        extraData[1];
} __attribute__((packed));

/*
CPU Disasm
Address   Hex dump          Command
00000000  90            NOP
00000001  90            NOP
00000002  68 EFBEADDE   PUSH DEADBEEF       Push address of DLL path
00000007  BA EFBEADDE   MOV EDX,DEADBEEF    move address of GetModuleHandle in EDX
0000000C  FFD2          CALL EDX            call it
0000000E  50            PUSH EAX            push return value of GetModuleHandle
0000000F  BA EFBEADDE   MOV EDX,DEADBEEF    grab the address of FreeLibrary
00000014  FFD2          CALL EDX            call it...
00000016  C20400        RET 0x04
00000019  00                                Filler
*/

// EB FE
#define EJECT_STUB      "\x90\x90"                \
                        "\x68\xEF\xBE\xAD\xDE"    \
                        "\xBA\xEF\xBE\xAD\xDE"    \
                        "\xFF\xD2"                \
                        "\x50"                    \
                        "\xBA\xEF\xBE\xAD\xDE"    \
                        "\xFF\xD2"                \
                        "\xC2\x04\x00"            \
                        "\x00"
struct EjectStub
{
    uint8_t     res1_[3];
    uintptr_t   dllString;
    uint8_t     res2_;
    uintptr_t   getModuleHandle;
    uint8_t     res3_[4];
    uintptr_t   freeLibrary;
    uint8_t     res5_[6];
    char        extraData[1];
} __attribute__((packed));


void Usage(char *exename)
{
    fprintf(stderr, "Usage: %s <command> <params>\n\n"
                    "|  Command  |      Params                      |\n"
                    "+-----------+----------------------------------+\n"
                    "| inject    | <pid> <dll path>                 |\n"
                    "| eject     | <pid> <dll name>                 |\n"
                    "| call      | <pid> <dll name> <function name> |\n", exename);
}

// Verify arguments and return the correct command ID, or exit
int VerifyArgs(int argc, char **argv)
{

    if (argc == 4 && strcmp(argv[1], "inject") == 0)
        return DO_INJECT;
    if (argc == 4 && strcmp(argv[1], "eject") == 0)
        return DO_EJECT;
    if (argc == 5 && strcmp(argv[1], "call") == 0)
        return DO_CALL;

    Usage(argv[0]);
    exit(1);

    return -1;
}

HANDLE InjectorOpenProcess(char *pid)
{
    printf("[+] Opening process with pid %s", pid);
    fflush(stdout);
    HANDLE proc = OpenProcess(PROCESS_INJECT_DLL, 0, atol(pid));
    if (proc == NULL)
    {
        printf("          [ failed  ]\n");
        exit(1);
    }

    printf("          [ success ]\n");
    return proc;
}

void *InjectorAlloc(HANDLE proc, DWORD size, DWORD protection)
{
    printf("[+] Allocating a memory page");
    fflush(stdout);
    void *page = VirtualAllocEx(proc, NULL, size, MEM_COMMIT, protection);
    if (page == NULL)
    {
        printf("               [ failed  ]\n");
        exit(1);
    }

    printf("               [ success ]\n");
    return page;
}

void InjectorWPM(HANDLE proc, void *dest, void *source, DWORD size)
{
    printf("[+] Writing some memory");
    fflush(stdout);
    if (!WriteProcessMemory(proc, dest, source, size, NULL))
    {
        printf("                    [ failed  ]\n");
        exit(1);
    }
    printf("                    [ success ]\n");
}

void InjectorCreateThread(HANDLE proc, void *procedure, void *argument)
{
    printf("[+] Creating remote thread");
    fflush(stdout);
    if (!CreateRemoteThread(proc, NULL, 0, procedure, argument, 0, NULL))
    {
        printf("                 [ failed  ]\n");
        exit(1);
    }
    printf("                 [ success ]\n");
}

// Eject a dll from target process
int DllEject(char *pid, char *dllName)
{
    printf("Ejecting %s from %s\n", dllName, pid);
    HANDLE proc = InjectorOpenProcess(pid);
    int size = sizeof(struct EjectStub) + strlen(dllName) + 1;
    void *page = InjectorAlloc(proc, size, PAGE_EXECUTE_READWRITE);
    void *gmhAddress = GetProcAddress(GetModuleHandle("kernel32.dll"), "GetModuleHandleA");
    void *flAddress = GetProcAddress(GetModuleHandle("kernel32.dll"), "FreeLibrary");

    struct EjectStub *cs = malloc(size);
    memcpy(cs, EJECT_STUB, sizeof(struct EjectStub));
    strcpy(cs->extraData, dllName);

    cs->getModuleHandle = (uintptr_t)gmhAddress;
    cs->freeLibrary = (uintptr_t)flAddress;
    cs->dllString = (uintptr_t)page + offsetof(struct EjectStub, extraData);
    InjectorWPM(proc, page, cs, size);
    DWORD lol;
    VirtualProtectEx(proc, page, size, PAGE_EXECUTE_READ, &lol);
    InjectorCreateThread(proc, page, NULL);
    return 0;
}

// Inject a DLL into the target process
int DllInject(char *pid, char *dllPath)
{
    printf("Injecting %s into %s\n", dllPath, pid);
    HANDLE proc = InjectorOpenProcess(pid);
    void *page = InjectorAlloc(proc, strlen(dllPath) + 1, PAGE_READWRITE);
    InjectorWPM(proc, page, dllPath, strlen(dllPath) + 1);
    void *loadLibraryAddress = GetProcAddress(GetModuleHandle("kernel32.dll"), "LoadLibraryA");
    InjectorCreateThread(proc, loadLibraryAddress, page);
    return 0;
}

// Call a function in a DLL in the target process
int DllCall(char *pid, char *dllName, char *functionName)
{
    printf("Calling %s@%s in %s\n", functionName, dllName, pid);
    HANDLE proc = InjectorOpenProcess(pid);
    int size = sizeof(struct CallStub) + strlen(dllName) + 1 + strlen(functionName) + 1;
    void *page = InjectorAlloc(proc, size, PAGE_EXECUTE_READWRITE);
    void *gmhAddress = GetProcAddress(GetModuleHandle("kernel32.dll"), "GetModuleHandleA");
    void *gpaAddress = GetProcAddress(GetModuleHandle("kernel32.dll"), "GetProcAddress");

    struct CallStub *cs = malloc(size);
    memcpy(cs, CALL_STUB, sizeof(struct CallStub));
    strcpy(cs->extraData, dllName);
    strcpy(cs->extraData + strlen(dllName) + 1, functionName);

    cs->getModuleHandle = (uintptr_t)gmhAddress;
    cs->getProcAddress = (uintptr_t)gpaAddress;
    cs->dllString = (uintptr_t)page + offsetof(struct CallStub, extraData);
    cs->funcString = (uintptr_t)page + offsetof(struct CallStub, extraData) + strlen(dllName) + 1;
    InjectorWPM(proc, page, cs, size);
    DWORD lol;
    VirtualProtectEx(proc, page, size, PAGE_EXECUTE_READ, &lol);
    InjectorCreateThread(proc, page, NULL);
    return 0;
}

int main(int argc, char **argv)
{
    int command = VerifyArgs(argc, argv);
    switch(command)
    {
        case DO_INJECT:
            return DllInject(argv[2], argv[3]);
        case DO_EJECT:
            return DllEject(argv[2], argv[3]);
        case DO_CALL:
            return DllCall(argv[2], argv[3], argv[4]);
    }
    return 1;
}

